#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/3/21
# @Author  : geekhch
# @Email   : geekhch@qq.com
# @Desc    : 词向量加载封装类，提供词向量索引转换接口

from gensim.models import KeyedVectors
from numpy import random
from myPath import *
from utils import logger
import numpy as np

class KeyedVocab:
    '''
    this class should be singleton;
    this class encapsulate methods for the transition among word-index-vector
    load pre-trained word2vector model with KeyedVectors
    tips: here we add <pad> <ukn>
    '''
    def __init__(self, model_path):
        logger.info(f'load pre-trained vocab: {model_path}')

        self.model = KeyedVectors.load_word2vec_format(model_path, datatype = np.float32)
        self.dim = self.model.vector_size

        self.token = ['<ukn>', '<pad>']
        random.seed(1)
        self.model.add(self.token, random.randn(len(self.token), self.dim)/10)



    def get_vec(self, word):
        return self.model.get_vector(word)

    def get_embedding(self):
        '''return all word vectors as numpy array'''
        return self.model.vectors

    def get_word2id_dict(self):
        '''返回word2id字典'''
        return dict(zip(self.model.index2word, range(self.__len__())))

    def word2index(self, w):
        try:
            return self.model.vocab[w].index
        except:
            # logger.warning(f'unknow word:{w}')
            return self.model.vocab['<ukn>'].index

    def index2word(self, idx):
        if idx >= len(self.model.vocab):
            return self.token[idx - len(self.model.vocab)]
        return self.model.index2word[idx]

    def __len__(self):
        return len(self.model.vocab)

vocab = KeyedVocab(VOCAB_CHAR300)
# vocab = KeyedVocab(VOCAB_SOUGOU)

if __name__ == '__main__':
    print(len(vocab))
    print('vector(“川”)=', vocab.get_vec('川'))